import java.util.HashSet;
import java.util.Set;

/**
 * 1261. 在受污染的二叉树中查找元素
 * https://leetcode-cn.com/problems/find-elements-in-a-contaminated-binary-tree/
 */
public class Solutions_1261 {
    public static void main(String[] args) {
//        String str = "[-1 -1 -1 -1 -1]";  // root:{0, 1, 2, 3, 4}
        String str = "[-1 # -1 -1 # -1]";  // root:{0, #, 2, 5, #, 11}
        TreeNode root = MyTreeNodeUtils.deserialize(str);

        FindElements obj = new FindElements(root);
        boolean res1 = obj.find(1);  // output: true
        System.out.println(res1);
        boolean res2 = obj.find(3);  // output: true
        System.out.println(res2);
        boolean res3 = obj.find(5);  // output: false
        System.out.println(res3);
    }
}

class FindElements {
    Set<Integer> set = null;
    public FindElements(TreeNode root) {
        set = new HashSet<>();
//        dfs(root, 0);
        root.val = 0;
        dfs1(root);
    }

    // 通过参数赋值
    void dfs(TreeNode root, int val) {
        if (root == null) {
            return;
        }
        root.val = val;
        set.add(val);
        // 左子节点值 = 根节点值 * 2 + 1，右子节点值 = 根节点值 * 2 + 2
        dfs(root.left, val * 2 + 1);
        dfs(root.right, val * 2 + 2);
    }

    // 通过父节点为子节点赋值
    // 左子节点值 = 根节点值 * 2 + 1，右子节点值 = 根节点值 * 2 + 2
    void dfs1(TreeNode root) {
        if (root == null) {
            return;
        }
        // 节点值记录到 Set 中，方便查找
        set.add(root.val);
        if (root.left != null) {
            // 为左子节点赋值，并递归
            root.left.val = root.val * 2 + 1;
            dfs1(root.left);
        }
        if (root.right != null) {
            // 为右子节点赋值，并递归
            root.right.val = root.val * 2 + 2;
            dfs1(root.right);
        }
    }

    public boolean find(int target) {
        return set.contains(target);
    }
}
